import cv2
import numpy as np
import matplotlib.pyplot as plt
# 参数对比
# 使用不同的以下参数进行对比。
# d: 9 - 滤波器的直径。更大的值意味着更大的邻域，但计算量更大。
# sigmaColor: 75 - 颜色空间标准差。值越大，更多的颜色会被混合（平滑），意味着更多不相似的像素会参与平滑。
# sigmaSpace: 75 - 空间坐标标准差。值越大，越远的像素也会被考虑在内。
# (5, 5): 高斯核的大小，必须是奇数。


def add_gaussian_noise(image, mean=0, sigma=25):
    """
    给图像添加高斯噪声。
    """
    row, col, ch = image.shape
    gauss = np.random.normal(mean, sigma, (row, col, ch))
    noisy_image = image + gauss
    noisy_image = np.clip(noisy_image, 0, 255) # 将像素值裁剪到[0, 255]范围
    return noisy_image.astype(np.uint8)

def _save_result(original_image, noisy_image, bilateral_filtered_image, gaussian_filtered_image, d, sigmaColor, sigmaSpace, kernel_size):
    plt.figure(figsize=(15, 10))

    plt.subplot(2, 2, 1)
    plt.imshow(original_image)
    plt.title('Original Image')
    plt.axis('off')

    plt.subplot(2, 2, 2)
    plt.imshow(noisy_image)
    plt.title('Noisy Image')
    plt.axis('off')

    plt.subplot(2, 2, 3)
    plt.imshow(bilateral_filtered_image)
    plt.title(f'Bilateral Filtered (d={d}, sigmaColor={sigmaColor}, sigmaSpace={sigmaSpace})')
    plt.axis('off')

    plt.subplot(2, 2, 4)
    plt.imshow(gaussian_filtered_image)
    plt.title(f'Gaussian Filtered (kernel={kernel_size})')
    plt.axis('off')

    plt.tight_layout()
    plt.savefig(f'./bilateral_filter/bilateral_d{d}_sigmaC{sigmaColor}_sigmaS{sigmaSpace}_gaussian_{kernel_size[0]}.png')
    plt.close()

def _filter_images(noisy_image, bilateral_params, gaussian_kernel_sizes):
    for d, sigmaColor, sigmaSpace in bilateral_params:
        bilateral_filtered_image = cv2.bilateralFilter(noisy_image, d=d, sigmaColor=sigmaColor, sigmaSpace=sigmaSpace)
        for kernel_size in gaussian_kernel_sizes:
            gaussian_filtered_image = cv2.GaussianBlur(noisy_image, kernel_size, 0)
            _save_result(original_image, noisy_image, bilateral_filtered_image, gaussian_filtered_image, d, sigmaColor, sigmaSpace, kernel_size)


def plot_images(original_image, noisy_image, bilateral_filtered_image, gaussian_filtered_image):
    plt.figure(figsize=(15, 10))

    plt.subplot(2, 2, 1)
    plt.imshow(original_image)
    plt.title('Original Image')
    plt.axis('off')

    plt.subplot(2, 2, 2)
    plt.imshow(noisy_image)
    plt.title('Noisy Image')
    plt.axis('off')

    plt.subplot(2, 2, 3)
    plt.imshow(bilateral_filtered_image)
    plt.title('Bilateral Filtered Image (Edges Preserved)')
    plt.axis('off')

    plt.subplot(2, 2, 4)
    plt.imshow(gaussian_filtered_image)
    plt.title('Gaussian Filtered Image (Edges Blurred)')
    plt.axis('off')

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 1. 加载一张图片
    image_path = './bilateral_filter/test_image.jpg'
    original_image = cv2.imread(image_path)
    original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB) # OpenCV默认读取BGR，转为RGB方便matplotlib显示

    # 2. 模拟添加一些高斯噪声
    noisy_image = add_gaussian_noise(original_image, mean=0, sigma=25)

    # 定义双边滤波参数
    bilateral_params = [
        (3, 120, 30),
        (9, 120, 30),
        (15, 120, 30),
        (3, 240, 30),
        (9, 240, 30),
        (15, 240, 30),
        # (9, 30, 30),
        # (9, 75, 75),
        # (15, 150, 150)
    ]

    # 定义高斯滤波核大小
    gaussian_kernel_sizes = [
        (5, 5),
    ]

    _filter_images(noisy_image, bilateral_params, gaussian_kernel_sizes)